import os
import sqlite3
import json
import pandas as pd
from datetime import datetime

# --- 配置 ---
PROMOTION_DIR = "推广"  # 推广文件夹相对于脚本的位置
DATABASE_PATH = "promotion_data.sqlite" # SQLite数据库路径 (修改到主文件夹)
CLEAR_EXISTING_DATA = False # 是否在每次导入前清空旧数据
# ------------

def check_and_upgrade_database(conn):
    """检查数据库结构，需要进行完全重建"""
    cursor = conn.cursor()
    
    # 检查是否存在旧版表
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND (name='Files' OR name='DataRows');")
    old_tables = cursor.fetchall()
    
    if old_tables:
        print("检测到旧版数据库结构，需要进行重建...")
        # 备份旧数据
        try:
            cursor.execute("CREATE TABLE IF NOT EXISTS old_files AS SELECT * FROM Files;")
            cursor.execute("CREATE TABLE IF NOT EXISTS old_datarows AS SELECT * FROM DataRows;")
            print("已备份旧数据到old_files和old_datarows表")
        except:
            print("备份旧数据失败，可能是首次结构调整")
        
        # 删除旧表
        cursor.execute("DROP TABLE IF EXISTS Files;")
        cursor.execute("DROP TABLE IF EXISTS DataRows;")
        print("已删除旧表结构")
        
        conn.commit()
        return True
    return False

def create_database_schema(conn):
    """创建新的数据库表结构"""
    cursor = conn.cursor()
    
    # 创建Categories表（替代原Files表）
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS Categories (
        category TEXT PRIMARY KEY,      -- 类别名称（文件名，不含路径和后缀）
        file_path TEXT NOT NULL,        -- 原始文件路径
        sheet_name TEXT NOT NULL,       -- 工作表名
        headers TEXT NOT NULL,          -- JSON array string of headers
        import_timestamp DATETIME NOT NULL
    );
    """)
    
    # 创建DataItems表（替代原DataRows表）
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS DataItems (
        item_id INTEGER PRIMARY KEY AUTOINCREMENT,
        category TEXT NOT NULL,         -- 关联到Categories表的类别
        row_index INTEGER NOT NULL,     -- 行号 (从0开始)
        row_data TEXT NOT NULL,         -- JSON object string of the row
        FOREIGN KEY (category) REFERENCES Categories (category)
    );
    """)
    
    # 为DataItems.category创建索引以提高查询效率
    cursor.execute("CREATE INDEX IF NOT EXISTS idx_dataitems_category ON DataItems (category);")
    
    conn.commit()
    print("数据库表结构创建完毕。")

def clear_data(conn):
    """清空数据"""
    cursor = conn.cursor()
    cursor.execute("DELETE FROM DataItems;")
    cursor.execute("DELETE FROM Categories;")
    # 重置自增ID
    cursor.execute("DELETE FROM sqlite_sequence WHERE name='DataItems';")
    conn.commit()
    print("数据已清空。")

def get_category_from_filename(file_path):
    """从文件路径中提取类别名称（文件名，不含路径和后缀）"""
    # 提取文件名
    file_name = os.path.basename(file_path)
    # 去除扩展名
    category = os.path.splitext(file_name)[0]
    # 去除可能的空格
    category = category.strip()
    return category

def process_excel_file(file_path, conn):
    """处理单个Excel文件"""
    print(f"正在处理文件: {file_path} ...")
    try:
        # 使用pandas读取Excel文件，读取所有sheet
        xls = pd.ExcelFile(file_path)
        cursor = conn.cursor()
        
        # 获取类别名称
        category = get_category_from_filename(file_path)
        
        for sheet_name in xls.sheet_names:
            print(f"  处理工作表: {sheet_name}")
            df = pd.read_excel(xls, sheet_name=sheet_name)
            # 处理空文件或空Sheet的情况
            if df.empty:
                print(f"  工作表 {sheet_name} 为空，跳过。")
                continue
                
            # 获取表头 (列名)
            headers = [str(col) for col in df.columns] # 确保表头是字符串
            headers_json = json.dumps(headers, ensure_ascii=False)

            # 将类别信息插入Categories表
            import_time = datetime.now()
            rel_file_path = os.path.relpath(file_path, os.path.dirname(__file__)) # 获取相对路径
            
            # 先检查是否已存在该类别
            cursor.execute("SELECT category FROM Categories WHERE category = ?", (category,))
            exists = cursor.fetchone()
            
            if not exists:
                # 插入新类别
                cursor.execute("""
                INSERT INTO Categories (category, file_path, sheet_name, headers, import_timestamp)
                VALUES (?, ?, ?, ?, ?)
                """, (category, rel_file_path, sheet_name, headers_json, import_time))
            else:
                # 更新现有类别
                cursor.execute("""
                UPDATE Categories SET 
                file_path = ?, sheet_name = ?, headers = ?, import_timestamp = ?
                WHERE category = ?
                """, (rel_file_path, sheet_name, headers_json, import_time, category))
            
            # 删除该类别的现有数据
            cursor.execute("DELETE FROM DataItems WHERE category = ?", (category,))
            # 逐行处理数据
            rows_to_insert = []
            for index, row in df.iterrows():
                # 将行数据转换为字典，键是表头，值是单元格数据
                row_dict = {}
                for header in headers:
                   value = row[header]
                   # JSON不支持NaN, NaT等特殊类型, 转换为None或字符串
                   if pd.isna(value):
                       row_dict[header] = None
                   elif isinstance(value, (pd.Timestamp, datetime)):
                        row_dict[header] = value.isoformat() # 转为ISO格式字符串
                   else:
                       row_dict[header] = value
                
                row_data_json = json.dumps(row_dict, ensure_ascii=False, default=str) # 添加 default=str 以防有其他无法序列化的类型
                rows_to_insert.append((category, index, row_data_json))
            # 批量插入行数据
            if rows_to_insert:
                cursor.executemany("""
                INSERT INTO DataItems (category, row_index, row_data)
                VALUES (?, ?, ?)
                """, rows_to_insert)
                print(f"    成功导入 {len(rows_to_insert)} 行数据。")
        conn.commit()
    except FileNotFoundError:
        print(f"错误：文件未找到 {file_path}")
    except Exception as e:
        conn.rollback() # 出错时回滚事务
        print(f"处理文件 {file_path} 时发生错误: {e}")


def main():
    """主函数"""
        
    # 确保 推广 目录存在
    if not os.path.isdir(PROMOTION_DIR):
        print(f"错误：推广目录 '{PROMOTION_DIR}' 不存在或不是一个目录。请确保脚本相对于项目根目录运行，并且推广目录存在。")
        return

    # 连接数据库 (如果不存在则会创建)
    conn = sqlite3.connect(DATABASE_PATH)
    print(f"已连接到数据库: {DATABASE_PATH}")

    try:
        # 检查并升级数据库结构
        need_rebuild = check_and_upgrade_database(conn)
        
        # 创建新表结构
        create_database_schema(conn)

        # 清空旧数据 (如果配置了)
        if CLEAR_EXISTING_DATA or need_rebuild:
            clear_data(conn)

        # 遍历推广目录下的所有 .xlsx 文件
        for filename in os.listdir(PROMOTION_DIR):
            if filename.lower().endswith(".xlsx"):
                file_path = os.path.join(PROMOTION_DIR, filename)
                process_excel_file(file_path, conn)

        print("\n数据导入完成！")

    except sqlite3.Error as e:
        print(f"数据库操作发生错误: {e}")
    finally:
        # 关闭数据库连接
        if conn:
            conn.close()
            print("数据库连接已关闭。")

if __name__ == "__main__":
    main() 